import torch
import tqdm
from torch import nn
import data_store
import gxl_config


class ChatModel(nn.Module):
    def __init__(self, vocab_size, embedding_size=80, num_hiddens=256, num_layers=2, dropout=0.0852):
        super(ChatModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.rnn = nn.GRU(embedding_size, num_hiddens, num_layers, dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def forward(self, X):
        X = self.embedding(X)
        X = X.permute(1, 0, 2)
        output, state = self.rnn(X)
        output = self.dense(output).permute(1, 0, 2)
        return output


def sequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项"""
    maxlen = X.size(1)
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value
    return X


# @save
class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""

    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    def forward(self, pred, label, valid_len):
        weights = torch.ones_like(label)
        weights = sequence_mask(weights, valid_len)
        self.reduction = 'none'
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(
            pred.permute(0, 2, 1), label)  # 规定需要(batch_size,vocab_size,step_nums)
        weighted_loss = (unweighted_loss * weights).mean(dim=1)
        return weighted_loss


def train(epochs, lr, model, train_iter, device):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss = MaskedSoftmaxCELoss()
    model.to(device)
    tqdm_object = tqdm.tqdm(range(epochs))
    for epoch in tqdm_object:
        model.train()
        loss_epoch = 0
        batch_size = 0
        for batch in train_iter:
            optimizer.zero_grad()
            X, X_valid_len, Y, Y_valid_len = [x.to(device) for x in batch]
            output = model(X)
            # print(output.shape, Y.shape, Y_valid_len.shape)
            l = loss(output, Y, Y_valid_len)
            l.sum().backward()
            optimizer.step()
            loss_epoch += l.mean()
            batch_size += 1
        tqdm_object.set_description(f'loss:{loss_epoch / batch_size}')
        torch.save(model.state_dict(),
                   gxl_config.MODEL_SAVE_PATH + 'gxl_chat_novel' + str(epoch) + gxl_config.get_now() + '.pt')


def run():
    # save_vocab()
    batch_size, epochs, lr = 32, 10, 0.01
    loader, vocab = data_store.getDataLoader(gxl_config.DATA2_FILE_PATH,batch_size)
    model = ChatModel(vocab_size=len(vocab), embedding_size=80, num_hiddens=256, num_layers=2, dropout=0.1)
    model.load_state_dict(torch.load(gxl_config.MODEL_SAVE_PATH + 'gxl_chat_novel92023-08-22_14:56:34.pt'))
    device = torch.device('cuda:2')
    train(epochs, lr, model, loader, device)


def predict(sentence="老师好", res_len=300):
    vocab = data_store.EasyVocab.load_data_to_get_vocab()
    model = ChatModel(vocab_size=len(vocab), embedding_size=80, num_hiddens=256, num_layers=2, dropout=0.1)
    model.load_state_dict(torch.load(gxl_config.MODEL_SAVE_PATH + 'gxl_chat_novel92023-08-22_14:56:34.pt'))
    model.to(torch.device('cpu'))
    sentence = [vocab[i] for i in sentence]
    sentence = torch.tensor(sentence, dtype=torch.long)
    sentence = sentence.unsqueeze(0)
    sentence = sentence.to(torch.device('cpu'))
    res = []
    while (True):
        output = model(sentence)
        output_tokens = output.argmax(dim=2)
        token_res = output_tokens[0, -1].item()
        if len(res) > 10 and (
                token_res == vocab['。'] or token_res == vocab['？'] or token_res == vocab['！']) or len(
            res) == res_len:
            break
        res.append(token_res)
        sentence = output_tokens
    return vocab.decode(res)


def save_vocab():
    vocab = data_store.EasyVocab()
    vocab.add_corpus(open(gxl_config.DATA1_FILE_PATH, 'r', encoding='utf8'))
    vocab.add_corpus(open(gxl_config.DATA2_FILE_PATH, 'r', encoding='utf8'))
    print('save vocab successfully, vocab length: {}'.format(len(vocab)))


if __name__ == '__main__':
    # run()
    # save_vocab()
    print(predict())
